import numpy as np
import pandas as pd
import os.path
import os
import time
import tensorflow as tf
# import tensorflow_addons as tfa
from midas import Midas

runs = 1

path = 'adult/data_tmp/'

def waiting(path):
    # dots = "."
  while not os.path.exists(path+'marker.csv'):
    if os.path.exists(path+'marker.csv'):
      break
    # if dots == "......":
    #   dots = "      "
    # elif dots == "      ":
    #   dots = "."
    # else:
    #   dots += "."
    time.sleep(0.5)
# print("\rWaiting"+dots)
  print("\rInit...       ")


def reverse_dummies(imputation, cat_cols):
  target_cols = imputation[cat_cols]
  imputation.drop(cat_cols, axis=1, inplace= True)
  imputation['y_1'] = target_cols.idxmax(axis=1)
  return imputation


def write_imputations(path, imputer):#, shape, nominal_cols= None):

  imputations = imputer.generate_samples(m=10).output_list
  n = 1
  for i in imputations:
    file_out = path + "midas_imp_" + str(n) + ".csv"
    # if nominal_cols is not None:
    # i = reverse_dummies(i, nominal_cols)
    # i = i.mul(tx_std).add(tx_mean)
    # assert i.shape[0] == shape
    i.to_csv(file_out, index=False)
    n +=1

  time.sleep(2)
  os.remove(path+'marker.csv')


categorical = ['workclass','marital_status','relationship',
               'race','class_labels','sex']

while not os.path.exists(path+'marker.csv'):

  if runs > 300:
  	break

  waiting(path)

  data_0 = pd.read_csv(path+"adult_midas_data.csv")

  data_0.columns.str.strip()

  data_1 = data_0[categorical]
  data_0.drop(categorical, axis= 1, inplace= True)

  constructor_list = [data_0]
  columns_list = []

  for column in data_1.columns:
    na_temp = data_1[column].isnull()
    temp = pd.get_dummies(data_1[column], prefix = column)
    temp[na_temp] = np.nan
    constructor_list.append(temp)
    columns_list.append(list(temp.columns.values))

  data_0 = pd.concat(constructor_list, axis=1)

  na_loc = data_0.isnull()
  data_0[na_loc] = np.nan

  imputer = Midas(layer_structure= [256,256], vae_layer= False, seed= 89, input_drop = 0.75)
  imputer.build_model(data_0, softmax_columns= columns_list)
  imputer.train_model(training_epochs = 20)

  write_imputations(path, imputer)

  runs += 1
